#include"StringPool.h"
#include"variable.h"
#include<vector>
#include<ctime>
#include<iostream>
#include<map>
#include<iomanip>
#include"DependencyInstance.h"
#include"Util.h"
#include"KBestParseForest2OSibling.h"
using namespace std;

KBestParseForest2OSibling::KBestParseForest2OSibling(int start, int end, DependencyInstance &inst, int K){
	this->start = start;
	this->end = end;
	this->K = K;
	this->sent = inst.forms;
	this->pos = inst.postags;
	this->chart = new ParseForestItem[(end + 1) * (end + 1) * 2 * 3 * K];

	for(int i = 0; i< end + 1; i++){
		for(int j = 0; j < end + 1; j++){
			for(int k = 0; k < 2; k++){
				for(int l = 0; l < 3; l++){
					for(int m = 0; m < K; m++){
						int key = getKey(i, j, k, l, m);
						chart[key] = ParseForestItem(i, -1, j, -1, k, l, Negative_Infinity, NULL, NULL);
					}
				}
			}
		}
	}
}

vector<pair<FeatureVector*, string*> >* KBestParseForest2OSibling::getBestParses(DependencyInstance *inst, DependencyPipe *pipe){
	vector<pair<FeatureVector*, string*> >* d = new vector<pair<FeatureVector*, string*> >(K);
	for(int i = 0; i < K; i++){
		int key = getKey(0, end, 0, 1, i);
		if(chart[key].prob > Negative_Infinity){
			(*d)[i].second = getDepString(chart[key]);
			(*d)[i].first = getFeatureVector(inst, pipe, (*d)[i].second);
		}
		else{
			(*d)[i].first = NULL;
			(*d)[i].second = NULL;
		}
	}
	return d;
}

vector<pair<FeatureVector*, string*> >* KBestParseForest2OSibling::getBestParses(DependencyInstance *inst, DependencyPipe *pipe, std::vector<vector<vector<double> > > &probs, std::vector<vector<vector<double> > > &probs_trips, std::vector<vector<vector<double> > > &probs_sibs, std::vector<vector<vector<vector<double> > > > &nt_probs, bool labeled){
	vector<pair<FeatureVector*, string*> >* d = new vector<pair<FeatureVector*, string*> >(K);
	vector<ParseForestItem> result = vector<ParseForestItem>(K);
	for(int i = 0; i < K; i++){
		result[i] = ParseForestItem(0, -1, end, -1, 0, 1, Negative_Infinity, NULL, NULL);
	}
	for(int i = 1; i < end + 1; i++){
		ParseForestItem *left = getItems(1, i, 1, 1);
		ParseForestItem *right = getItems(i, end, 0, 1);
		vector<BestPair>* pairs = getKBestPairs(left, right);
		int size = (int)(pairs->size());
		bool added = true;
		double prob_0i = probs[0][i][0];
		double nt_prob_s_01 = 0.0;
		double nt_prob_t_00 = 0.0;
		if(labeled){
			nt_prob_s_01 = nt_probs[0][rootType][0][1];
			nt_prob_t_00 = nt_probs[i][rootType][0][0];
			prob_0i += nt_prob_s_01 + nt_prob_t_00;
		}

		double prob_00i = probs_trips[0][0][i] + probs_sibs[0][i][0];

		for(int k = 0; k < size; k++){
			if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
				break;
			}

			int k1 = (*pairs)[k].first;
			int k2 = (*pairs)[k].second;
			double prob = left[k1].prob + right[k2].prob + prob_0i + prob_00i;
			added = false;
			for(int j = 0; j < K; j++){
				if(prob > result[j].prob){
					ParseForestItem tmp = result[j];
					result[j] = ParseForestItem(0, i, end, -1, 0, 1, prob, &(left[k1]), &(right[k2]));
					for(int m = j + 1; m < K; m++){
						ParseForestItem tmp1 = result[m];
						result[m] = tmp;
						tmp = tmp1;
					}
					added = true;
					break;
				}
			}
			if(!added){
				break;
			}
		}
		delete(pairs);
	}
	for(int i = 0; i < K; i++){
		if(result[i].prob > Negative_Infinity){
			string* str1 = getDepString(*(result[i].left));
			(*d)[i].second = Util::trim(str1);
			pool->push_back(str1);
			int ll = (int)((*d)[i].second->length());
			if(ll != 0){
				(*d)[i].second->append(" ");
			}
			(*d)[i].second->append("0|");
			str1 = Util::intToString(result[i].r);
			str1->append(":");
			(*d)[i].second->append(*str1);
			pool->push_back(str1);
			str1 = Util::intToString(KBestParseForest::rootType);
			str1->append(" ");
			(*d)[i].second->append(*str1);
			pool->push_back(str1);
			str1 = getDepString(*(result[i].right));
			string* str2 = Util::trim(str1);
			pool->push_back(str1);
			(*d)[i].second->append(*str2);
			pool->push_back(str2);
			str1 = (*d)[i].second;
			(*d)[i].second = Util::trim(str1);
			pool->push_back(str1);

			(*d)[i].first = getFeatureVector(inst, pipe, (*d)[i].second);
		}
		else{
			(*d)[i].first = NULL;
			(*d)[i].second = NULL;
		}
	}
	return d;
}

int KBestParseForest2OSibling::getKey(int s, int t, int dir, int comp, int k){
	int key = s;
	key = key * (end + 1) + t;
	key = key * 2 + dir;
	key = key * 3 + comp;
	key = key * K + k;
	return key;
}

string* KBestParseForest2OSibling::getDepString(ParseForestItem &pfi){
	if(pfi.left == NULL){
		return pool->get();
	}

	if(pfi.dir == 0 && pfi.comp == 0){
		string* str1 = getDepString(*(pfi.left));
		str1->append(" ");
		string* str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		string* result = Util::trim(str1);
		pool->push_back(str1);
		pool->push_back(str2);
		result->append(" ");
		str1 = Util::intToString(pfi.s);
		result->append(*str1 + "|");
		pool->push_back(str1);
		str1 = Util::intToString(pfi.t);
		result->append(*str1 + ":");
		pool->push_back(str1);
		str1 = Util::intToString(pfi.type);
		result->append(*str1);
		pool->push_back(str1);
		str1 = result;
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else if(pfi.dir == 1 && pfi.comp == 0){
		string* str1 = Util::intToString(pfi.t);
		str1->append("|");
		string* str2 = Util::intToString(pfi.s);
		str2->append(":");
		str1->append(*str2);
		pool->push_back(str2);
		str2 = Util::intToString(pfi.type);
		str2->append(" ");
		str1->append(*str2);
		pool->push_back(str2);
		string* result = str1;
		str1 = getDepString(*(pfi.left));
		str1->append(" ");
		str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		pool->push_back(str2);
		str2 = Util::trim(str1);
		pool->push_back(str1);
		result->append(*str2);
		pool->push_back(str2);
		str1 = result;
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else{
		string* str1 = getDepString(*(pfi.left));
		str1->append(" ");
		string* str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		string* result = Util::trim(str1);
		pool->push_back(str1);
		pool->push_back(str2);
		return result;
	}
}


KBestParseForest2OSibling::~KBestParseForest2OSibling(){}
